# src/influxdb_client.py
from influxdb_client import InfluxDBClient, Point
from influxdb_client.client.write_api import SYNCHRONOUS, ASYNCHRONOUS
import pandas as pd
import logging
from typing import List, Dict
from datetime import datetime
from logger import logger
import os


class InfluxDBManager:
    def __init__(self, url: str, token: str, org: str, bucket: str):
        self.client = InfluxDBClient(url=url, token=token, org=org, timeout=30*1000)
        self.write_api = self.client.write_api(write_options=ASYNCHRONOUS)
        self.query_api = self.client.query_api()
        self.bucket = bucket
        self.org = org
        self.logger = logger

        # 测试连接
        self.test_connection()

    def test_connection(self):
        """测试连接"""
        try:
            ready = self.client.ping()
            if ready:
                self.logger.info("InfluxDB 连接成功")
            else:
                self.logger.error("InfluxDB 连接失败")
        except Exception as e:
            self.logger.error(f"InfluxDB 连接测试失败: {e}")

    def write_klines_data(self, df: pd.DataFrame):
        """
        写入K线数据到InfluxDB
        """
        if df.empty:
            self.logger.warning("没有数据可写入")
            return

        points = []

        for _, row in df.iterrows():
            point = Point("klines") \
                .tag("symbol", row['symbol']) \
                .tag("interval", row['interval']) \
                .field("open", float(row['open'])) \
                .field("high", float(row['high'])) \
                .field("low", float(row['low'])) \
                .field("close", float(row['close'])) \
                .field("volume", float(row['volume'])) \
                .field("quote_asset_volume", float(row['quote_asset_volume'])) \
                .field("number_of_trades", int(row['number_of_trades'])) \
                .field("taker_buy_base_volume", float(row['taker_buy_base_asset_volume'])) \
                .field("taker_buy_quote_volume", float(row['taker_buy_quote_asset_volume'])) \
                .time(row['open_time'])

            points.append(point)

        try:
            self.write_api.write(bucket=self.bucket, record=points)
            self.logger.info(f"成功写入 {len(points)} 条K线数据")
        except Exception as e:
            self.logger.error(f"写入数据失败: {e}")

    def write_current_price(self, symbol: str, price: float):
        """
        写入当前价格
        """
        point = Point("current_price") \
            .tag("symbol", symbol) \
            .field("price", price) \
            .time(datetime.utcnow())

        try:
            self.write_api.write(bucket=self.bucket, record=point)
            self.logger.debug(f"写入当前价格: {symbol} = {price}")
        except Exception as e:
            self.logger.error(f"写入当前价格失败: {e}")

    def query_klines_data(self, symbol: str, interval: str, start: str = "-7d", end: str = "now()") -> pd.DataFrame:
        """
        查询K线数据
        """
        query = f'''
        from(bucket: "{self.bucket}")
            |> range(start: {start}, stop: {end})
            |> filter(fn: (r) => r._measurement == "klines")
            |> filter(fn: (r) => r.symbol == "{symbol}")
            |> filter(fn: (r) => r.interval == "{interval}")
            |> pivot(rowKey:["_time"], columnKey: ["_field"], valueColumn: "_value")
        '''

        try:
            result = self.query_api.query_data_frame(query)
            if not result.empty:
                # 重命名列
                result = result.rename(columns={'_time': 'open_time'})
                return result
            else:
                return pd.DataFrame()
        except Exception as e:
            self.logger.error(f"查询数据失败: {e}")
            return pd.DataFrame()

    def get_latest_timestamp(self, symbol: str, interval: str) -> datetime:
        """
        获取最新数据的时间戳
        """
        query = f'''
        from(bucket: "{self.bucket}")
            |> range(start: -30d)
            |> filter(fn: (r) => r._measurement == "klines")
            |> filter(fn: (r) => r.symbol == "{symbol}")
            |> filter(fn: (r) => r.interval == "{interval}")
            |> last()
        '''

        try:
            result = self.query_api.query(query)
            if result and len(result) > 0:
                for table in result:
                    for record in table.records:
                        return record.get_time()
            return None
        except Exception as e:
            self.logger.error(f"获取最新时间戳失败: {e}")
            return None

    def close(self):
        """关闭连接"""
        self.write_api.close()
        self.client.close()

def main(frequency):
    path = f"G:/binance_k_line/v1/{frequency}"
    files = os.listdir(path)
    print(files)
    count = 0
    exist_symbols = get_exist_symbols(frequency)
    symbols = ["ETHUSDT", "ETHBTC", "BNBUSDT", "BNBBTC", "XRPUSDT", "XRPBTC", "SOLUSDT", "TRXUSDT", "DOGEUSDT", "SOLBTC", "TRXBTC", "DOGEBTC"]
    for file in files:
        # if file != "arbzilla.csv":
        #     continue
        # print(file)
        count += 1
        print(f"{count}/{len(files)}")
        symbol = file.split("_")[1]
        if symbol in exist_symbols:
            print(f"{symbol} is exist")
            continue
        if symbol not in symbols:
            continue
        print(f"处理{symbol}的数据")
        df = pd.read_csv(os.path.join(path, file))
if __name__ == '__main__':
    my_token = "wikfeUowyxSF5nQsNaO37k6BcI89b8wyqk4vayJ-9Y3I5p2kQdrUMiMUGD7ciCjw4zAQMdowzNwzPAKuWp7qqQ=="
    # my_token = "YzTN8RQIPgFluNXeUkmY3pq_PTNdgh6sPEEdArIP2YhbdzdpUnN9bxFmBh4ZvJOIfTnYb-PhA4lkfWUECvUePw=="
    # client = InfluxDBClient(url="http://localhost:8086", token=my_token, org="tianzhi", timeout=30*1000)
    # ready = client.ping()
    # print(ready)
    manager = InfluxDBManager(url="http://localhost:8086", token=my_token,
                              org="tianzhi", bucket="binance_data")
    # manager.test_connection()